package com.hksj.common.core.executor;

import cn.hutool.core.collection.CollectionUtil;
import lombok.NonNull;
import org.slf4j.MDC;
import org.springframework.lang.Nullable;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;

import java.util.Map;
import java.util.Objects;
import java.util.concurrent.Callable;
import java.util.concurrent.Future;
import java.util.concurrent.RejectedExecutionHandler;

/**
 * @author fjt
 * @data 2024/8/8 16:56
 * @describe 自定义线程池
 */
public class MyThreadPoolTaskExecutor extends ThreadPoolTaskExecutor {

	/**
	 * 线程池配置
	 */
	public MyThreadPoolTaskExecutor() {

	}

	/**
	 * 线程池配置
	 *
	 * @param corePoolSize             核心线程数
	 * @param maxPoolSize              最大线程数
	 * @param keepAliveTime            存活时长
	 * @param queueCapacity            队列长度
	 * @param poolNamePrefix           线程前缀
	 * @param rejectedExecutionHandler 线程拒绝策略
	 */
	public MyThreadPoolTaskExecutor(@Nullable int corePoolSize,
									@Nullable int maxPoolSize,
									@Nullable int keepAliveTime,
									@Nullable int queueCapacity,
									String poolNamePrefix,
									@Nullable RejectedExecutionHandler rejectedExecutionHandler) {
		setCorePoolSize(corePoolSize);
		setMaxPoolSize(maxPoolSize);
		setKeepAliveSeconds(keepAliveTime);
		setQueueCapacity(queueCapacity);
		setThreadNamePrefix(poolNamePrefix);
		setRejectedExecutionHandler(rejectedExecutionHandler);
		initialize();
	}

	private Map<String, String> getContextForTask() {
		return MDC.getCopyOfContextMap();
	}


	@Override
	public void execute(@NonNull Runnable command) {
		super.execute(wrap(command, getContextForTask()));
	}

	@NonNull
	@Override
	public Future<?> submit(@NonNull Runnable task) {
		return super.submit(wrap(task, getContextForTask()));
	}

	@NonNull
	@Override
	public <T> Future<T> submit(@NonNull Callable<T> task) {
		return super.submit(wrap(task, getContextForTask()));
	}

	private static <T> Callable<T> wrap(final Callable<T> task, final Map<String, String> context) {
		return () -> {
			if (CollectionUtil.isNotEmpty(context)) {
				MDC.setContextMap(context);
			}

			try {
				return task.call();
			} finally{
				if (CollectionUtil.isNotEmpty(context) && Objects.nonNull(MDC.getMDCAdapter())) {
					MDC.clear();
				}
			}
		};
	}

	private static Runnable wrap(final Runnable runnable, final Map<String, String> context) {
		return () -> {
			if (CollectionUtil.isNotEmpty(context)) {
				MDC.setContextMap(context);
			}

			try {
				runnable.run();
			} finally {
				if (CollectionUtil.isNotEmpty(context) && Objects.nonNull(MDC.getMDCAdapter())) {
					MDC.clear();
				}
			}
		};
	}

}
